use spin::Lazy;
use x86_64::{
    registers::control::Cr2,
    set_general_handler,
    structures::idt::{InterruptDescriptorTable, InterruptStackFrame, PageFaultErrorCode},
    VirtAddr,
};

use crate::arch::{context::THREAD_MANAGER, smp::CPUS};

use super::gdt::DOUBLE_FAULT_IST_INDEX;

pub static IDT: Lazy<InterruptDescriptorTable> = Lazy::new(|| {
    let mut idt = InterruptDescriptorTable::new();

    set_general_handler!(&mut idt, unexcepted_interrupt, 0..255);

    idt.breakpoint.set_handler_fn(breakpoint);
    idt.segment_not_present.set_handler_fn(segment_not_present);
    idt.invalid_opcode.set_handler_fn(invalid_opcode);
    idt.page_fault.set_handler_fn(page_fault);
    idt.general_protection_fault
        .set_handler_fn(general_protection_fault);

    unsafe {
        idt.double_fault
            .set_handler_fn(double_fault)
            .set_stack_index(DOUBLE_FAULT_IST_INDEX as u16);
    }

    idt[InterruptIndex::Timer as u8].set_handler_fn(timer_interrupt);
    idt[InterruptIndex::ApicError as u8].set_handler_fn(lapic_error);
    idt[InterruptIndex::ApicSpurious as u8].set_handler_fn(spurious_interrupt);
    idt[InterruptIndex::HpetTimer as u8].set_handler_fn(hpet_interrupt);

    idt
});

fn unexcepted_interrupt(stack: InterruptStackFrame, index: u8, error_code: Option<u64>) {
    panic!(
        "Unexcepted interrupt\nstack = {:#?}, index = {}, error_code = {:#?}",
        stack, index, error_code
    );
}

fn timer_handler(last_context: VirtAddr) -> VirtAddr {
    super::apic::end_of_interrupt();
    THREAD_MANAGER.lock().schedule(last_context)
}

extern "x86-interrupt" fn lapic_error(_frame: InterruptStackFrame) {
    log::error!("Local APIC error!");
    super::apic::end_of_interrupt();
}

extern "x86-interrupt" fn spurious_interrupt(_frame: InterruptStackFrame) {
    log::debug!("Received spurious interrupt!");
    super::apic::end_of_interrupt();
}

#[naked]
extern "x86-interrupt" fn timer_interrupt(_frame: InterruptStackFrame) {
    unsafe {
        core::arch::naked_asm!(
            "cli",
            crate::push_context!(),
            "mov rdi, rsp",
            "call {timer_handler}",
            "mov rsp, rax",
            crate::pop_context!(),
            "sti",
            "iretq",
            timer_handler = sym timer_handler,
        );
    }
}

extern "x86-interrupt" fn hpet_interrupt(_frame: InterruptStackFrame) {}

extern "x86-interrupt" fn segment_not_present(frame: InterruptStackFrame, error_code: u64) {
    error!("Exception: Segment Not Present\n{:#?}", frame);
    error!("Error Code: {:#x}", error_code);
    error!("Unrecoverable fault occured, halting!");
    loop {
        x86_64::instructions::hlt();
    }
}

extern "x86-interrupt" fn general_protection_fault(frame: InterruptStackFrame, error_code: u64) {
    error!("Exception: General Protection Fault\n{:#?}", frame);
    error!("Error Code: {:#x}", error_code);
    loop {
        x86_64::instructions::hlt();
    }
}

extern "x86-interrupt" fn invalid_opcode(frame: InterruptStackFrame) {
    error!("Exception: Invalid Opcode\n{:#?}", frame);
    loop {
        x86_64::instructions::hlt();
    }
}

extern "x86-interrupt" fn breakpoint(frame: InterruptStackFrame) {
    debug!("Exception: Breakpoint\n{:#?}", frame);
}

extern "x86-interrupt" fn double_fault(frame: InterruptStackFrame, error_code: u64) -> ! {
    error!("Exception: Double Fault\n{:#?}", frame);
    error!("Error Code: {:#x}", error_code);
    x86_64::instructions::interrupts::disable();
    loop {
        x86_64::instructions::hlt();
    }
}

extern "x86-interrupt" fn page_fault(frame: InterruptStackFrame, error_code: PageFaultErrorCode) {
    warn!("Exception: Page Fault\n{:#?}", frame);
    warn!("Error Code: {:#x}", error_code);
    match Cr2::read() {
        Ok(address) => {
            warn!("Fault Address: {:#x}", address);
        }
        Err(error) => {
            warn!("Invalid virtual address: {:?}", error);
        }
    }
    loop {
        x86_64::instructions::hlt();
    }
}

const INTERRUPT_INDEX_OFFSET: u8 = 32;

#[derive(Debug, Clone, Copy)]
#[repr(u8)]
pub enum InterruptIndex {
    Timer = INTERRUPT_INDEX_OFFSET,
    ApicError,
    ApicSpurious,
    HpetTimer,
    Keyboard,
    Mouse,
}

pub fn init() {
    IDT.load();
    super::apic::init();
    CPUS.write().init_ap();
}
